import sys
parent_dir = "../pygl"
sys.path.append(parent_dir)
import ctypes
import OpenGL.GL as gl
import glm
import imgui
import glfw
from glfw import _GLFWwindow as GLFWwindow
import numpy as np
import pygl
from common import App, ScreenQuad

class Mandelbrot(App):
    # 初始化着色器，vao vbo等
    def setupVertex(self) -> None:
        self.program = pygl.ProgramCompute("intermediate/shaders/mandelbrot.comp")
        self.render_texture = pygl.RenderTexture2D(1, gl.GL_RGBA32F, self.width, self.height)
        self.render_texture.bingImage(0, 0, gl.GL_WRITE_ONLY)
        self.quad = ScreenQuad(self.render_texture)
        self.num_groups = [0] * 3
        self.num_groups[0] = int((self.width + self.program.local_size[0] - 1) / self.program.local_size[0])
        self.num_groups[1] = int((self.height + self.program.local_size[1] - 1) / self.program.local_size[1])
        self.num_groups[2] = 1
        self.scale = 2.0 / float(self.height)
        self.pan_x = self.width * 0.75
        self.pan_y = self.height * 0.5
        self.last_pos = glm.vec2()
        self.center_pos = glm.vec2()
        self.max_iter = 100
        self.imgui_window_width = 320
        self.imgui_window_height = 180
        
    # 配置imgui界面元素
    def showImGui(self)->None:
        imgui.begin("设置")
        if imgui.button("重置"):
            self.scale = 2.0 / float(self.height)
            self.pan_x = self.width * 0.75
            self.pan_y = self.height * 0.5
        imgui.text("按住鼠标右键拖动可以移动画面")
        imgui.text("滚动鼠标滚轮可以缩放画面")
        imgui.text("向前滚动是放大；\n向后滚动是缩小。")
        _, self.max_iter = imgui.slider_int(
            "最大迭代次数", self.max_iter, 50, 1000)
        imgui.end()       

    # 渲染图形
    def render(self) -> None:
        self.program.use()
        self.program.setUniform2f(0, self.pan_x, self.pan_y)
        self.program.setUniform1f(1, self.scale)
        self.program.setUniform1i(2, self.max_iter)
        self.program.dispatch(self.num_groups[0], self.num_groups[1], self.num_groups[2])
        gl.glMemoryBarrier(gl.GL_SHADER_IMAGE_ACCESS_BARRIER_BIT)
        self.quad.render(None, None)

    # 释放资源
    def cleanup(self) -> None:
        self.program.delete()
        self.render_texture.delete()
        self.quad.delete()

    def mouse_move(self, window:GLFWwindow, x:float, y:float)->None:
        if imgui.is_any_item_active():
            return
        deltaX = x - self.last_pos.x
        deltaY = y - self.last_pos.y
        if glfw.get_mouse_button(self.window , glfw.MOUSE_BUTTON_RIGHT) == glfw.PRESS:
            self.pan_x += deltaX
            self.pan_y -= deltaY
        self.last_pos = glm.vec2(x, y)

    def mouse_wheel(self, window:GLFWwindow, xoffset:float, yoffset:float)-> None:
        self.center_pos = glm.vec2(xoffset,
                                self.height - yoffset)
        self.zoom(-yoffset * 10.0)

    def mouse_click(self, window:GLFWwindow, button:int, action:int, 
                    mods:int)->None:
        xpos, ypos = glfw.get_cursor_pos(self.window)
        self.last_pos = glm.vec2(xpos, ypos)
        self.center_pos = glm.vec2(self.last_pos.x, 
                            self.height - self.last_pos.y)

    def zoom(self, delta: float) -> None:
        cx = self.scale * (self.center_pos.x - self.pan_x)
        cy = self.scale * (self.center_pos.y - self.pan_y)
        self.scale *= pow(1.01, delta)
        self.pan_x = (self.scale * self.center_pos.x - cx) / self.scale
        self.pan_y = (self.scale * self.center_pos.y - cy) / self.scale

    def resize(self, window:GLFWwindow, width:int, height:int)-> None:
        self.width = width
        self.height = height
        gl.glViewport(0, 0, width, height)
        self.render_texture.resize(width, height)
        self.render_texture.bingImage(0, 0, gl.GL_WRITE_ONLY)
        self.num_groups[0] = int((self.width + self.program.local_size[0] - 1) / self.program.local_size[0])
        self.num_groups[1] = int((self.height + self.program.local_size[1] - 1) / self.program.local_size[1])

if __name__ == "__main__":
    app = Mandelbrot(100, 100, 1000, 600, "你好，曼德尔布诺特App！")
    app.run()


